from __future__ import division
import numpy as np
#from numpy import multiply as nm
from Constant_Physics import Physics
from Tools import *
from MRTools import *
from MRTools import Lorentzian as lorentz
import matplotlib.pyplot as plt
import src_fortran.hamiltonians as hamiltonians
import src_fortran.kronecker as kronecker
from inspect import stack #To "automatize the clk functions"

class Shallow_Donor(Physics):

    def __init__(self, clk, specie="P",fortran_compiled=True):
        self.clk=clk
        self.fortran_compiled_functions=fortran_compiled
        if specie == "P":
            self.nspin = 1 / 2
            self.dim_n=2
            self.A = 117.53    #MHz
            self.ge = 1.9985
            self.gn = -2.2632
        elif specie == "As":
            self.nspin = 3 / 2
            self.A = 198.35   #MHz
            self.ge = 1.9984
            self.gn = -0.762601
            self.dim_n=4
        elif specie == "Sb5/2":
            self.nspin = 5 / 2
            self.A = 186.802   #MHz
            self.ge = 1.9986
            self.gn = -1.3454
            self.dim_n=6
        elif specie == "Sb7/2":
            self.nspin = 7 / 2
            self.A = 101.516    #MHz
            self.ge = 1.9986
            self.gn = -0.728492
            self.dim_n=8
        elif specie == "Bi":
            self.nspin = 9 / 2
            self.A = 1475.4    #MHz
            self.ge = 2.0003
            self.gn = -0.914
            self.dim_n=10

        self.dim_e=2
        self.espin=1/2
        self.dimension = self.dim_e* self.dim_n
        self.hamiltonian = np.zeros([self.dimension, self.dimension])
        self.transformation_matrix = np.zeros([self.dimension, self.dimension])
        
        self.wn=Physics.EnergyUnit * (10 ** (-6)) / Physics.h * self.gn * Physics.mun
        self.we=Physics.EnergyUnit * (10 ** (-6)) / Physics.h * self.ge * Physics.mue


        if self.fortran_compiled_functions == False:
            #Preparation of the nuclear spin matrices
            self.Ip = np.zeros([self.nspin * 2 + 1, self.nspin * 2 + 1])
            list_jm = np.arange(-self.nspin, self.nspin + 1)
            for n in range(int(self.nspin * 2)):
                self.Ip[n , n + 1] = np.sqrt((self.nspin - list_jm[n]) * (self.nspin + list_jm[n] + 1))
            self.Ix = (1 / 2) * (self.Ip + np.transpose(self.Ip))
            self.Iy = (1 / 2) * (np.transpose(self.Ip) - self.Ip)
            self.Iz = np.diag(-list_jm)
            self.I = [self.Ix, self.Iy, self.Iz]
            self.Id = np.diag(np.ones(self.nspin * 2 + 1))
    
            self.espin = 1 / 2  #Electronic spin
            self.dim_e=2
            self.Sx = 1 / 2 * np.array([[0, 1], [ 1, 0]])
            self.Sy = 1 / 2 * np.array([[0, -1], [1, 0]])
            self.Sz = 1 / 2 * np.array([[1, 0], [0, -1]])
            self.S = [self.Sx, self.Sy, self.Sz]
            self.Sd = np.diag([1, 1])
        elif self.fortran_compiled_functions==True:
            self.Sd=hamiltonians.build_sd(2)
            self.Sx=hamiltonians.build_sx(2)
            self.Sy=hamiltonians.build_sy(2)
            self.Sz=hamiltonians.build_sz(2)
            self.Id=hamiltonians.build_sd(self.dim_n)
            self.Ix=hamiltonians.build_sx(self.dim_n)
            self.Iy=hamiltonians.build_sy(self.dim_n)
            self.Iz=hamiltonians.build_sz(self.dim_n)


    def D_Hamiltonian_Zn(self, b):
        """
        Input:
            -B: the magnetic field [Bx,By,Bz] in tesla
        Output:
            -hamiltonian: the donor hamiltonian (Zn)
        """
        if self.fortran_compiled_functions==True:
            return kronecker.sptv(self.dim_n,self.Ix,self.Iy,self.Iz,b,self.wn)    #Zn term
        else:
            return self.wn*(self.Ix*b[0]+self.Iy*b[1]+self.Iz*b[2])
    
    def D_Hamiltonian_Ze(self, b):
        """
        Input:
            -B: the magnetic field [Bx,By,Bz]
        Output:
            -hamiltonian: the donor hamiltonian (Ze)
        """
        if self.fortran_compiled_functions==True:
            return kronecker.sptv(self.dim_e,self.Sx,self.Sy,self.Sz,b,self.we)    #Ze term
        else:
            return self.we*(self.Sx*b[0]+self.Sy*b[1]+self.Sz*b[2])
        
    def D_Hamiltonian_hyperfine(self):
        """
        Output:
            -hyperfine
        """
        if self.fortran_compiled_functions==True:
            d=np.array([2,2,1,2,2,1])
            return kronecker.sp(self.dim_e,self.dim_n,1,-1,1,d,self.Sx,self.Sy,self.Sz,self.Ix,self.Iy,self.Iz,self.A) 
        else:
            return self.A*(np.kron(self.Sx,self.Ix)-np.kron(self.Sy,self.Iy)+np.kron(self.Sz,self.Iz))


    def D_Hamiltonian_Dz(self, b):
        """
        Input:
            -B: the magnetic field [Bx,By,Bz]
        Output:
            -hamiltonian: the donor hamiltonian (Ze+Zn+H) 
        """
        if self.fortran_compiled_functions==True:
            d=np.array([2,2,1,2,2,1])
            res=kronecker.k22(self.D_Hamiltonian_Ze(b),self.dim_e,self.dim_e,1,self.Id,self.dim_n,self.dim_n,1,1) 
            res+=kronecker.k22(self.Sd,self.dim_e,self.dim_e,1,self.D_Hamiltonian_Zn(b),self.dim_n,self.dim_n,1,1)
            res+=self.D_Hamiltonian_hyperfine()
            return res
        else:
            return np.kron(self.D_Hamiltonian_Ze(b),self.Id) + np.kron(self.Sd,self.D_Hamiltonian_Zn(b))+ self.D_Hamiltonian_hyperfine() 

    def F_sweep_simu(self, field, transition="EPR"):
        """
        Checked
        """
        width = 10 #MHz
        n_points = 10000
        Y = np.zeros(n_points)
        X = np.arange(0, n_points, 1)
        print field
        ham = self.D_Hamiltonian_Dz(field.field_vec())
        eigvalDz, eigvecDz = np.linalg.eig(ham)  #Diagonalization
        print eigvalDz 
        P = eigvecDz
        iP = np.linalg.inv(P)
        if transition == "EPR":
            MR = np.kron(self.Sx, self.Id)
        elif transition == "NMR":
            MR = np.kron(self.Sd, self.Ix)
        MR_Mb = np.dot(iP, np.dot(MR, P))

        lorentzian = lorentz()
        for i in range(MR_Mb[:, 0].size):
            for j in range(MR_Mb[0, :].size):
                center = np.abs(eigvalDz[i] - eigvalDz[j])
                coeffs = MR_Mb[i, j] ** 2, center, width
                evalpts = X
                Y += lorentzian.Eval(coeffs, evalpts)
                #print center,MR_Mb[i, j] ** 2

        #t1 = np.arange(0, 10, 0.1)
        #t2 = lorentzian.Eval([1, 4, 1], t1)


        plt.plot(X, Y, '-')
        plt.show()

if __name__ == "__main__":
    
    class SDR_parameters:
        #SDR model parameters
        PE = 1e8
        Cp = 1e4
        Cap = 1e4 
        Rp = 2e3
        Rap = 1e5
        
        #Modulation
        modulation_type="field" #field/AM/FM
        modulation_frequency=100000 #Hz
        modulation_phase=0 #rad
        modulation_shape="sin" #sin/square
        modulation_amplitude=0.001 #T/ percent/ MHz
    
        #Signal properties        
        linewidth_MHz=1
        linewidth_mT=1
        
        #Experimental conditions
        temperature=16#K 
    
    ##########################################
    ##########################################
    class Simulation_parameters:
        time_range={}
        time_range['without resonance']=1e-4
        time_range['with resonance']=1e-4
        time_intervale={}
        time_intervale['without resonance']=1e-7
        time_intervale['with resonance']=1e-7
        integration_time=1e-4#s, 10 periods
        transient_std_length=5# in tau unit
        #do_adjust_time: If yes, fit the given transient with an exponential, 
        #find the associated tau, and extend the simulation time
        do_adjust_time=False 
    
    
    ##########################################
    ##########################################
    class Bfield(Physics,SDR_parameters):
        
        def __init__(self,clk,B0=0.1,B0xyz=np.array([0, 0, 1]),modulated=True,fortran_compiled=True):
            self.modulated=modulated
            self.clk=clk
            self.fortran_compiled=fortran_compiled
            self.B0=B0
            self.B0xyz=B0xyz
        
        def field_vec(self,t=0):
            """
            Input:
                -t: the time
            Output:
                -field vector
            """
            res=np.array([0,0,0])
            if self.fortran_compiled==True:
                if self.modulated:
                    res= hamiltonians.bfield(self.B0*Physics.FieldUnit,
                                     self.B0xyz,
                                     SDR_parameters.modulation_frequency,
                                     SDR_parameters.modulation_phase,
                                     SDR_parameters.modulation_amplitude*Physics.FieldUnit,
                                     t)
                else:
                    res= hamiltonians.bfield(self.B0*Physics.FieldUnit,
                                     self.B0xyz,
                                     1,
                                     0,
                                     0,
                                     0)
            else:
                if self.modulated:
                    res = (self.B0 + SDR_parameters.modulation_amplitude*Physics.FieldUnit* np.sin(2*np.pi*t*SDR_parameters.modulation_frequency+SDR_parameters.modulation_phase)) / np.linalg.norm(self.B0xyz) * self.B0xyz
                else:
                    res = (self.B0 ) / np.linalg.norm(self.B0xyz) * self.B0xyz
                
            return res
            
        def field_norm(self,t=0):
            """
            Input:
                -t: the time
            Output:
                -field norm
            """
            self.clk.TIC(self.__class__.__name__+":"+stack()[0][3],"L")
            if self.modulated:
                return (self.B0 + SDR_parameters.modulation_amplitude* np.sin(2*np.pi*t*SDR_parameters.modulation_frequency+SDR_parameters.modulation_phase))
            else:
                return self.B0
            self.clk.TOC(self.__class__.__name__+":"+stack()[0][3])

    clk=Computation_time()
    Bi = Shallow_Donor(clk,"Bi")
    Bvec=Bfield(clk,0.066056)
    Bi.F_sweep_simu(Bvec)
    #print Bi.Iy


